#define _GNU_SOURCE 1
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
// #include <stropts.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <termios.h>
#include <unistd.h>

static void fail_io(const char *msg, ...)
    __attribute__((noreturn))
    __attribute__((format(printf, 1, 2)));

/* Prints MSG, formatting as with printf(),
   plus an error message based on errno,
   and exits. */
static void fail_io(const char *msg, ...)
{
    va_list args;

    va_start(args, msg);
    vfprintf(stderr, msg, args);
    va_end(args);

    if (errno != 0)
        fprintf(stderr, ": %s", strerror(errno));
    putc('\n', stderr);
    exit(EXIT_FAILURE);
}

/* If FD is a terminal, configures it for noncanonical input mode
   with VMIN and VTIME set as indicated.
   If FD is not a terminal, has no effect. */
static void make_noncanon(int fd, int vmin, int vtime)
{
    if (isatty(fd))
    {
        struct termios termios;
        if (tcgetattr(fd, &termios) < 0)
            fail_io("tcgetattr");
        termios.c_lflag &= ~(ICANON | ECHO);
        termios.c_cc[VMIN] = vmin;
        termios.c_cc[VTIME] = vtime;
        if (tcsetattr(fd, TCSANOW, &termios) < 0)
            fail_io("tcsetattr");
    }
}

/* Make FD non-blocking if NONBLOCKING is true,
   or blocking if NONBLOCKING is false. */
static void make_nonblocking(int fd, bool nonblocking)
{
    int flags = fcntl(fd, F_GETFL);
    if (flags < 0)
        fail_io("fcntl");
    if (nonblocking)
        flags |= O_NONBLOCK;
    else
        flags &= ~O_NONBLOCK;
    if (fcntl(fd, F_SETFL, flags) < 0)
        fail_io("fcntl");
}

/* Handle a read or write on *FD, which is the pty if FD_IS_PTY
   is true, that returned end-of-file or error indication RETVAL.
   The system call is named CALL, for use in error messages.
   Sets *FD to -1 if the fd is no longer readable or writable. */
static void handle_error(ssize_t retval, int *fd, bool fd_is_pty, const char *call)
{
    if (fd_is_pty)
    {
        if (retval < 0)
        {
            if (errno == EIO)
            {
                /* Slave side of pty has been closed. */
                *fd = -1;
            }
            else
                fail_io(call);
        }
    }
    else
    {
        if (retval == 0)
        {
            close(*fd);
            *fd = -1;
        }
        else
            fail_io(call);
    }
}

/* Copies data from stdin to PTY and from PTY to stdout until no
   more data can be read or written. */
static void relay(int pty, int dead_child_fd)
{
    struct pipe
    {
        int in, out;
        char buf[BUFSIZ];
        size_t size, ofs;
        bool active;
    };
    struct pipe pipes[2];

    /* Make PTY, stdin, and stdout non-blocking. */
    make_nonblocking(pty, true);
    make_nonblocking(STDIN_FILENO, true);
    make_nonblocking(STDOUT_FILENO, true);

    /* Configure noncanonical mode on PTY and stdin to avoid
       waiting for end-of-line.  We want to minimize context
       switching on PTY (for efficiency) and minimize latency on
       stdin to avoid a laggy user experience. */
    make_noncanon(pty, 16, 1);
    make_noncanon(STDIN_FILENO, 1, 0);

    memset(pipes, 0, sizeof pipes);
    pipes[0].in = STDIN_FILENO;
    pipes[0].out = pty;
    pipes[1].in = pty;
    pipes[1].out = STDOUT_FILENO;

    while (pipes[1].in != -1)
    {
        fd_set read_fds, write_fds;
        int retval;
        int i;

        FD_ZERO(&read_fds);
        FD_ZERO(&write_fds);
        for (i = 0; i < 2; i++)
        {
            struct pipe *p = &pipes[i];

            /* Don't do anything with the stdin->pty pipe until we
               have some data for the pty->stdout pipe.  If we get
               too eager, Bochs will throw away our input. */
            if (i == 0 && !pipes[1].active)
                continue;

            if (p->in != -1 && p->size + p->ofs < sizeof p->buf)
                FD_SET(p->in, &read_fds);
            if (p->out != -1 && p->size > 0)
                FD_SET(p->out, &write_fds);
        }
        FD_SET(dead_child_fd, &read_fds);

        do
        {
            retval = select(FD_SETSIZE, &read_fds, &write_fds, NULL, NULL);
        } while (retval < 0 && errno == EINTR);
        if (retval < 0)
            fail_io("select");

        if (FD_ISSET(dead_child_fd, &read_fds))
            break;

        for (i = 0; i < 2; i++)
        {
            struct pipe *p = &pipes[i];
            if (p->in != -1 && FD_ISSET(p->in, &read_fds))
            {
                ssize_t n = read(p->in, p->buf + p->ofs + p->size,
                                 sizeof p->buf - p->ofs - p->size);
                if (n > 0)
                {
                    p->active = true;
                    p->size += n;
                    if (p->size == BUFSIZ && p->ofs != 0)
                    {
                        memmove(p->buf, p->buf + p->ofs, p->size);
                        p->ofs = 0;
                    }
                }
                else
                    handle_error(n, &p->in, p->in == pty, "read");
            }
            if (p->out != -1 && FD_ISSET(p->out, &write_fds))
            {
                ssize_t n = write(p->out, p->buf + p->ofs, p->size);
                if (n > 0)
                {
                    p->ofs += n;
                    p->size -= n;
                    if (p->size == 0)
                        p->ofs = 0;
                }
                else
                    handle_error(n, &p->out, p->out == pty, "write");
            }
        }
    }

    if (pipes[1].out == -1)
        return;

    make_nonblocking(STDOUT_FILENO, false);
    for (;;)
    {
        struct pipe *p = &pipes[1];
        ssize_t n;

        /* Write buffer. */
        while (p->size > 0)
        {
            n = write(p->out, p->buf + p->ofs, p->size);
            if (n < 0)
                fail_io("write");
            else if (n == 0)
                fail_io("zero-length write");
            p->ofs += n;
            p->size -= n;
        }
        p->ofs = 0;

        p->size = n = read(p->in, p->buf, sizeof p->buf);
        if (n <= 0)
            return;
    }
}

static int dead_child_fd;

static void sigchld_handler(int signo __attribute__((unused)))
{
    if (write(dead_child_fd, "", 1) < 0)
        _exit(1);
}

int main(int argc __attribute__((unused)), char *argv[])
{
    int master, slave;
    char *name;
    pid_t pid;
    struct sigaction sa;
    int pipe_fds[2];
    struct itimerval zero_itimerval, old_itimerval;

    if (argc < 2)
    {
        fprintf(stderr,
                "usage: squish-pty COMMAND [ARG]...\n"
                "Squishes both stdin and stdout into a single pseudoterminal,\n"
                "which is passed as stdout to run the specified COMMAND.\n");
        return EXIT_FAILURE;
    }

    /* Open master side of pty and get ready to open slave. */
    master = open("/dev/ptmx", O_RDWR | O_NOCTTY);
    if (master < 0)
        fail_io("open \"/dev/ptmx\"");
    if (grantpt(master) < 0)
        fail_io("grantpt");
    if (unlockpt(master) < 0)
        fail_io("unlockpt");

    /* Open slave side of pty. */
    name = ptsname(master);
    if (name == NULL)
        fail_io("ptsname");
    slave = open(name, O_RDWR);
    if (slave < 0)
        fail_io("open \"%s\"", name);

    /* System V implementations need STREAMS configuration for the
       slave. */
    // if (isastream (slave))
    //   {
    //     if (ioctl (slave, I_PUSH, "ptem") < 0
    //         || ioctl (slave, I_PUSH, "ldterm") < 0)
    //       fail_io ("ioctl");
    //   }

    /* Arrange to get notified when a child dies, by writing a byte
       to a pipe fd.  We really want to use pselect() and
       sigprocmask(), but Solaris 2.7 doesn't have it. */
    if (pipe(pipe_fds) < 0)
        fail_io("pipe");
    dead_child_fd = pipe_fds[1];

    memset(&sa, 0, sizeof sa);
    sa.sa_handler = sigchld_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART;
    if (sigaction(SIGCHLD, &sa, NULL) < 0)
        fail_io("sigaction");

    /* Save the virtual interval timer, which might have been set
       by the process that ran us.  It really should be applied to
       our child process. */
    memset(&zero_itimerval, 0, sizeof zero_itimerval);
    if (setitimer(ITIMER_VIRTUAL, &zero_itimerval, &old_itimerval) < 0)
        fail_io("setitimer");

    pid = fork();
    if (pid < 0)
        fail_io("fork");
    else if (pid != 0)
    {
        /* Running in parent process. */
        int status;
        close(slave);
        relay(master, pipe_fds[0]);

        /* If the subprocess has died, die in the same fashion.
           In particular, dying from SIGVTALRM tells the pintos
           script that we ran out of CPU time. */
        if (waitpid(pid, &status, WNOHANG) > 0)
        {
            if (WIFEXITED(status))
                return WEXITSTATUS(status);
            else if (WIFSIGNALED(status))
                raise(WTERMSIG(status));
        }
        return 0;
    }
    else
    {
        /* Running in child process. */
        if (setitimer(ITIMER_VIRTUAL, &old_itimerval, NULL) < 0)
            fail_io("setitimer");
        if (dup2(slave, STDOUT_FILENO) < 0)
            fail_io("dup2");
        if (close(pipe_fds[0]) < 0 || close(pipe_fds[1]) < 0 || close(slave) < 0 || close(master) < 0)
            fail_io("close");
        execvp(argv[1], argv + 1);
        fail_io("exec");
    }
}
